#include "segment_post_process.h"
#include "string.h"

#include "typeConv.hpp"
#include <algorithm>
#include <cmath>
#include <future>
#include <iomanip>
#include <iostream>
#include <opencv2/opencv.hpp>
#include <string>
#include <sys/time.h>
#include <utility>
#include <vector>
#ifdef LYNXI_PLUGIN
#include "lyn_plugin_dev.h"
#endif

// 模型配置信息
struct SegmentConfig {
  int class_num;
  float _classThreshold;
  float _nmsThreshold;
  float _maskThreshold;
  int _segChannels;
  std::vector<std::string> class_names;
};

// 配置类别
#define CLASS_NUM 80
const SegmentConfig default_segment_config = {
    CLASS_NUM,
    0.25,
    0.45,
    0.5,
    32,
    {"person",        "bicycle",      "car",
     "motorcycle",    "airplane",     "bus",
     "train",         "truck",        "boat",
     "traffic light", "fire hydrant", "stop sign",
     "parking meter", "bench",        "bird",
     "cat",           "dog",          "horse",
     "sheep",         "cow",          "elephant",
     "bear",          "zebra",        "giraffe",
     "backpack",      "umbrella",     "handbag",
     "tie",           "suitcase",     "frisbee",
     "skis",          "snowboard",    "sports ball",
     "kite",          "baseball bat", "baseball glove",
     "skateboard",    "surfboard",    "tennis racket",
     "bottle",        "wine glass",   "cup",
     "fork",          "knife",        "spoon",
     "bowl",          "banana",       "apple",
     "sandwich",      "orange",       "broccoli",
     "carrot",        "hot dog",      "pizza",
     "donut",         "cake",         "chair",
     "couch",         "potted plant", "bed",
     "dining table",  "toilet",       "tv",
     "laptop",        "mouse",        "remote",
     "keyboard",      "cell phone",   "microwave",
     "oven",          "toaster",      "sink",
     "refrigerator",  "book",         "clock",
     "vase",          "scissors",     "teddy bear",
     "hair drier",    "toothbrush"}};

template <typename T> T vectorProduct(const std::vector<T> &vec) {
  T product = 1;
  for (const auto &element : vec) {
    product *= element;
  }
  return product;
}

int lynSegmentPostProcess(SegmentPostProcessInfo_t *post_info) {
  using namespace std;
  using namespace cv;

  float *output_tensor = (float *)post_info->output_tensor;
  lynSegBoxesInfo *boxesInfo = post_info->boxesInfo;

  cv::Vec4d params{post_info->fatio, post_info->fatio, (double)post_info->padx,
                   (double)post_info->pady};

  int net_height = post_info->model_height;
  int net_width = post_info->model_width;

  int shapeTmp = net_height / 8 * net_width / 8;
  shapeTmp += net_height / 16 * net_width / 16;
  shapeTmp += net_height / 32 * net_width / 32;

  vector<int> out1Shape{1, CLASS_NUM + 4 + 32, shapeTmp};
  vector<int> out2Shape{1, 32, net_height / 4, net_width / 4};
  int seg_width = out2Shape[3];
  int seg_height = out2Shape[2];

  std::vector<cv::Mat> net_outputs;

  float *pp = output_tensor;
  Mat out1(out1Shape, CV_32F, pp);
  Mat out2(out2Shape, CV_32F, pp + vectorProduct(out1Shape));

  net_outputs.push_back(out1);
  net_outputs.push_back(out2);

  vector<int> class_ids;     // res-class_id
  vector<float> confidences; // res-conf
  vector<Rect> boxes;        // res-box
  vector<Mat> picked_proposals;

  int netOut_width = default_segment_config.class_names.size() + 4 +
                     default_segment_config._segChannels; // 4 + 80 + 32 = 116
  Mat output0 = Mat(Size(net_outputs[0].size[2], net_outputs[0].size[1]),
                    CV_32F, (float *)net_outputs[0].data)
                    .t(); //[bs,116,8400]=>[bs,8400,116]

  int rows = output0.rows; // 8400
  float *pdata = (float *)output0.data;
  for (int r = 0; r < rows; ++r) {
    Mat scores(1, default_segment_config.class_names.size(), CV_32FC1,
               pdata + 4);
    Point classIdPoint;
    double max_class_socre;
    minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
    max_class_socre = (float)max_class_socre;
    if (max_class_socre >= default_segment_config._classThreshold) {
      // vector<float> temp_proto(pdata + 4 +
      // default_segment_config.class_names.size(), pdata + netOut_width);
      // //mask 32
      cv::Mat temp_proto;
      int mask32Size =
          netOut_width - (4 + default_segment_config.class_names.size());
      temp_proto.create(1, mask32Size, CV_32FC1);
      memcpy(temp_proto.data,
             pdata + 4 + default_segment_config.class_names.size(),
             mask32Size * sizeof(float));
      picked_proposals.push_back(temp_proto);

      // rect [x,y,w,h]
      float x = (pdata[0] - params[2]) / params[0];
      float y = (pdata[1] - params[3]) / params[1];
      float w = pdata[2] / params[0];
      float h = pdata[3] / params[1];
      int left = MAX(int(x - 0.5 * w + 0.5), 0);
      int top = MAX(int(y - 0.5 * h + 0.5), 0);
      class_ids.push_back(classIdPoint.x);
      confidences.push_back(max_class_socre);
      boxes.push_back(Rect(left, top, int(w + 0.5), int(h + 0.5)));
    }
    pdata += netOut_width; // next line
  }
  // NMS
  vector<int> nms_result;
  dnn::NMSBoxes(boxes, confidences, default_segment_config._classThreshold,
                default_segment_config._nmsThreshold, nms_result);
  Rect holeImgRect(0, 0, post_info->img_width - 1, post_info->img_height - 1);

  boxesInfo->maskData = nullptr;
  boxesInfo->maskDataLen = 0;

  vector<std::shared_ptr<u_int8_t>> vecMaskDataPtr;
  boxesInfo->boxesNum = std::min((int)nms_result.size(), BOX_MAX_NUM);

  for (int i = 0; i < boxesInfo->boxesNum; i++) {
    int idx = nms_result[i];

    lynSegBox &segBox = boxesInfo->boxes[i];

    Rect tmpBox = boxes[idx] & holeImgRect;
    segBox.score = confidences[idx];
    segBox.xmin = tmpBox.x;
    segBox.ymin = tmpBox.y;
    segBox.width = tmpBox.width;
    segBox.height = tmpBox.height;
    strncpy(segBox.label,
            default_segment_config.class_names[class_ids[idx]].c_str(),
            TEXT_MAX_LEN - 1);

    // crop from mask_protos
    int rang_x =
        floor((tmpBox.x * params[0] + params[2]) / net_width * seg_width);
    int rang_y =
        floor((tmpBox.y * params[1] + params[3]) / net_height * seg_height);
    int rang_w = ceil(((tmpBox.x + tmpBox.width) * params[0] + params[2]) /
                      net_width * seg_width) -
                 rang_x;
    int rang_h = ceil(((tmpBox.y + tmpBox.height) * params[1] + params[3]) /
                      net_height * seg_height) -
                 rang_y;

    rang_w = MAX(rang_w, 1);
    rang_h = MAX(rang_h, 1);
    if (rang_x + rang_w > seg_width) {
      if (seg_width - rang_x > 0)
        rang_w = seg_width - rang_x;
      else
        rang_x -= 1;
    }
    if (rang_y + rang_h > seg_height) {
      if (seg_height - rang_y > 0)
        rang_h = seg_height - rang_y;
      else
        rang_y -= 1;
    }

    vector<Range> roi_rangs;
    roi_rangs.push_back(Range(0, 1));
    roi_rangs.push_back(Range::all());
    roi_rangs.push_back(Range(rang_y, rang_h + rang_y));
    roi_rangs.push_back(Range(rang_x, rang_w + rang_x));

    // crop
    Mat temp_mask_protos = net_outputs[1](roi_rangs).clone();
    Mat protos = temp_mask_protos.reshape(
        0, {default_segment_config._segChannels, rang_w * rang_h});
    Mat matmul_res = (picked_proposals[idx] * protos).t();
    Mat masks_feature = matmul_res.reshape(1, {rang_h, rang_w});
    Mat dest, mask;

    // sigmoid
    exp(-masks_feature, dest);
    dest = 1.0 / (1.0 + dest);

    int left = floor((net_width / seg_width * rang_x - params[2]) / params[0]);
    int top = floor((net_height / seg_height * rang_y - params[3]) / params[1]);
    int width = ceil(net_width / seg_width * rang_w / params[0]);
    int height = ceil(net_height / seg_height * rang_h / params[1]);

    cv::resize(dest, mask, Size(width, height), INTER_NEAREST);
    mask =
        mask(tmpBox - Point(left, top)) > default_segment_config._maskThreshold;

    segBox.maskCol = mask.cols;
    segBox.maskRow = mask.rows;
    segBox.id = class_ids[idx];
    boxesInfo->maskDataLen += mask.cols * mask.rows;

    std::shared_ptr<u_int8_t> maskPtr = std::shared_ptr<u_int8_t>(
        new u_int8_t[mask.cols * mask.rows], [](u_int8_t *p) { delete[] p; });
    std::memcpy(maskPtr.get(), mask.data, mask.total() * mask.elemSize());

    vecMaskDataPtr.push_back(maskPtr);
  }

  if (boxesInfo->maskDataLen > 0) {
    boxesInfo->maskData = new uint8_t[boxesInfo->maskDataLen];
    int pos = 0;
    for (int i = 0; i < boxesInfo->boxesNum; i++) {
      std::memcpy(boxesInfo->maskData + pos, vecMaskDataPtr[i].get(),
                  boxesInfo->boxes[i].maskCol * boxesInfo->boxes[i].maskRow);
      pos += boxesInfo->boxes[i].maskCol * boxesInfo->boxes[i].maskRow;
    }
  }

  return 0;
}

int deleteMaskData(lynSegBoxesInfo *boxesInfo) {
  delete[] boxesInfo->maskData;
  return 0;
}